/**
  * Copyright (c) CSIRO Australia, 2009
  *
  * @author $Author: jia020 $
  * @version $Id: XmlSignatureUtils.java 643 2009-05-20 00:20:55Z jia020 $
  */
package au.csiro.nt.pdsp.util;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.StringReader;
import java.io.StringWriter;
import java.security.InvalidAlgorithmParameterException;
import java.security.Key;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.UnrecoverableEntryException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;

import javax.xml.crypto.AlgorithmMethod;
import javax.xml.crypto.KeySelector;
import javax.xml.crypto.KeySelectorException;
import javax.xml.crypto.KeySelectorResult;
import javax.xml.crypto.MarshalException;
import javax.xml.crypto.XMLCryptoContext;
import javax.xml.crypto.XMLStructure;
import javax.xml.crypto.dsig.CanonicalizationMethod;
import javax.xml.crypto.dsig.DigestMethod;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.Transform;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureException;
import javax.xml.crypto.dsig.XMLSignatureFactory;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.dom.DOMValidateContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerConfigurationException;
import javax.xml.transform.TransformerException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import org.w3c.dom.Document;
import org.w3c.dom.NodeList;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;

public class XmlSignatureUtils {

	public static String signDoc(String theKeyStorePath, String theKeyStorePasswd, String theKeyName, String theInputFileName, String theOutputFileName)
	{
		// Create a DOM XMLSignatureFactory that will be used to
		// generate the enveloped signature.
		XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");

		// Create a Reference to the enveloped document (in this case,
		// you are signing the whole document, so a URI of "" signifies
		// that, and also specify the SHA1 digest algorithm and
		// the ENVELOPED Transform.
		Reference ref;
		try {
			ref = fac.newReference
			 ("", fac.newDigestMethod(DigestMethod.SHA1, null),
			  Collections.singletonList
			   (fac.newTransform
			    (Transform.ENVELOPED, (TransformParameterSpec) null)),
			     null, null);
			SignedInfo si = fac.newSignedInfo
			 (fac.newCanonicalizationMethod
			  (CanonicalizationMethod.INCLUSIVE,
			   (C14NMethodParameterSpec) null),
			    fac.newSignatureMethod(SignatureMethod.RSA_SHA1, null),
			     Collections.singletonList(ref));
			
			// Load the KeyStore and get the signing key and certificate.
			KeyStore ks = KeyStore.getInstance("JKS");
			ks.load(new FileInputStream(theKeyStorePath), theKeyStorePasswd.toCharArray());
			KeyStore.PrivateKeyEntry keyEntry =
			    (KeyStore.PrivateKeyEntry) ks.getEntry
			        (theKeyName, new KeyStore.PasswordProtection(theKeyStorePasswd.toCharArray()));
			X509Certificate cert = (X509Certificate) keyEntry.getCertificate();

			// Create the KeyInfo containing the X509Data.
			KeyInfoFactory kif = fac.getKeyInfoFactory();
			List x509Content = new ArrayList();
			x509Content.add(cert.getSubjectX500Principal().getName());
			x509Content.add(cert);
			X509Data xd = kif.newX509Data(x509Content);
			KeyInfo ki = kif.newKeyInfo(Collections.singletonList(xd));
						
			// Instantiate the document to be signed.
			DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
			dbf.setNamespaceAware(true);
			Document doc = dbf.newDocumentBuilder().parse
			    (new FileInputStream(theInputFileName));

			// Create a DOMSignContext and specify the RSA PrivateKey and
			// location of the resulting XMLSignature's parent element.
			DOMSignContext dsc = new DOMSignContext
			    (keyEntry.getPrivateKey(), doc.getDocumentElement());

			// Create the XMLSignature, but don't sign it yet.
			XMLSignature signature = fac.newXMLSignature(si, ki);

			// Marshal, generate, and sign the enveloped signature.
			signature.sign(dsc);
			
			// Output the resulting document.
			OutputStream os = new FileOutputStream(theOutputFileName);
			TransformerFactory tf = TransformerFactory.newInstance();
			Transformer trans = tf.newTransformer();
			trans.transform(new DOMSource(doc), new StreamResult(os));
			os.close();
			return theOutputFileName;
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (InvalidAlgorithmParameterException e) {
			e.printStackTrace();
		} catch (KeyStoreException e) {
			e.printStackTrace();
		} catch (CertificateException e) {
			e.printStackTrace();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (UnrecoverableEntryException e) {
			e.printStackTrace();
		} catch (SAXException e) {
			e.printStackTrace();
		} catch (ParserConfigurationException e) {
			e.printStackTrace();
		} catch (MarshalException e) {
			e.printStackTrace();
		} catch (XMLSignatureException e) {
			e.printStackTrace();
		} catch (TransformerConfigurationException e) {
			e.printStackTrace();
		} catch (TransformerException e) {
			e.printStackTrace();
		}
		return null;
	}
	
	public static String signContent(String theKeyStorePath, String theKeyStorePasswd, String theKeyName, String theContent)
	{
		// Create a DOM XMLSignatureFactory that will be used to
		// generate the enveloped signature.
		XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");

		// Create a Reference to the enveloped document (in this case,
		// you are signing the whole document, so a URI of "" signifies
		// that, and also specify the SHA1 digest algorithm and
		// the ENVELOPED Transform.
		Reference ref;
		try {
			ref = fac.newReference
			 ("", fac.newDigestMethod(DigestMethod.SHA1, null),
			  Collections.singletonList
			   (fac.newTransform
			    (Transform.ENVELOPED, (TransformParameterSpec) null)),
			     null, null);
			SignedInfo si = fac.newSignedInfo
			 (fac.newCanonicalizationMethod
			  (CanonicalizationMethod.INCLUSIVE,
			   (C14NMethodParameterSpec) null),
			    fac.newSignatureMethod(SignatureMethod.RSA_SHA1, null),
			     Collections.singletonList(ref));
			
			// Load the KeyStore and get the signing key and certificate.
			KeyStore ks = KeyStore.getInstance("JKS");
			ks.load(new FileInputStream(theKeyStorePath), theKeyStorePasswd.toCharArray());
			KeyStore.PrivateKeyEntry keyEntry =
			    (KeyStore.PrivateKeyEntry) ks.getEntry
			        (theKeyName, new KeyStore.PasswordProtection(theKeyStorePasswd.toCharArray()));
			X509Certificate cert = (X509Certificate) keyEntry.getCertificate();

			// Create the KeyInfo containing the X509Data.
			KeyInfoFactory kif = fac.getKeyInfoFactory();
			List x509Content = new ArrayList();
			x509Content.add(cert.getSubjectX500Principal().getName());
			x509Content.add(cert);
			X509Data xd = kif.newX509Data(x509Content);
			KeyInfo ki = kif.newKeyInfo(Collections.singletonList(xd));
						
			// Instantiate the document to be signed.
			DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
			dbf.setNamespaceAware(true);
			Document doc = dbf.newDocumentBuilder().parse
			    (new InputSource(new StringReader(theContent)));

			// Create a DOMSignContext and specify the RSA PrivateKey and
			// location of the resulting XMLSignature's parent element.
			DOMSignContext dsc = new DOMSignContext
			    (keyEntry.getPrivateKey(), doc.getDocumentElement());

			// Create the XMLSignature, but don't sign it yet.
			XMLSignature signature = fac.newXMLSignature(si, ki);

			// Marshal, generate, and sign the enveloped signature.
			signature.sign(dsc);
			
			// Output the resulting document.
			TransformerFactory tf = TransformerFactory.newInstance();
			Transformer trans = tf.newTransformer();
			StringWriter os = new StringWriter();
			trans.transform(new DOMSource(doc), new StreamResult(os));
			return os.toString();
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (InvalidAlgorithmParameterException e) {
			e.printStackTrace();
		} catch (KeyStoreException e) {
			e.printStackTrace();
		} catch (CertificateException e) {
			e.printStackTrace();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (UnrecoverableEntryException e) {
			e.printStackTrace();
		} catch (SAXException e) {
			e.printStackTrace();
		} catch (ParserConfigurationException e) {
			e.printStackTrace();
		} catch (MarshalException e) {
			e.printStackTrace();
		} catch (XMLSignatureException e) {
			e.printStackTrace();
		} catch (TransformerConfigurationException e) {
			e.printStackTrace();
		} catch (TransformerException e) {
			e.printStackTrace();
		}
		return null;
	}

	public static int validateDoc(Document theDoc) throws Exception
	{
		XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");
		// Find Signature element.
		NodeList nl =
		    theDoc.getElementsByTagNameNS(XMLSignature.XMLNS, "Signature");
		if (nl.getLength() == 0) {
		    throw new Exception("Cannot find Signature element");
		}

		// Create a DOMValidateContext and specify a KeySelector
		// and document context.
		DOMValidateContext valContext = new DOMValidateContext
		    (new X509KeySelector(), nl.item(0));

		// Unmarshal the XMLSignature.
		XMLSignature signature = fac.unmarshalXMLSignature(valContext);

		// Validate the XMLSignature.
		boolean coreValidity = signature.validate(valContext);
		// Check core validation status.
		if (coreValidity == false) {
		    System.err.println("Signature failed core validation");
		    boolean sv = signature.getSignatureValue().validate(valContext);
		    System.out.println("signature validation status: " + sv);
		    if (sv == false) {
		        // Check the validation status of each Reference.
		        Iterator i = signature.getSignedInfo().getReferences().iterator();
		        for (int j=0; i.hasNext(); j++) {
		            boolean refValid = ((Reference) i.next()).validate(valContext);
		            System.out.println("ref["+j+"] validity status: " + refValid);
		        }
		    }
		    return -1;
		} else {
		    System.out.println("Signature passed core validation");
		}
		return 0;
	}


	public static void main()
	{	
//            System.out.println(XmlSignatureUtils.signContent("C:\\home\\cwang\\dev\\workspace\\econtract\\conf\\partya.keystore",
//					"tester", 
//					"key-party-b", 
//					"<contract><participant>partya</participant><participant>partyb</participant></contract>"));
            String theFilePath="c:\\mykey.keystore";
            DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
            dbf.setNamespaceAware(true);
            Document doc;
            try
            {
                    doc = dbf.newDocumentBuilder().parse
                        (new File(theFilePath));
                    if (XmlSignatureUtils.validateDoc(doc) != 0)
                            System.out.println("failed to validate document "+theFilePath);
                    else
                            System.out.println("document "+theFilePath+" is valid!");
            } catch (SAXException e)
            {
                    e.printStackTrace();
            } catch (IOException e)
            {
                    e.printStackTrace();
            } catch (ParserConfigurationException e)
            {
                    e.printStackTrace();
            } catch (Exception e)
            {
                    e.printStackTrace();
            }     
        }
}
